package com.db_tenant;

import com.db_tenant.citus.CitusTenant;
import com.db_tenant.citus.CoordinateTenant;
import io.ebean.config.TenantDataSourceProvider;
import lombok.extern.slf4j.Slf4j;
import org.avaje.datasource.DataSourcePool;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.util.List;

@Slf4j
@Component
class CitusDataSrcProvider implements TenantDataSourceProvider {

  @Autowired
  private CoordinateTenant coordTenant;

  @Autowired
  @Qualifier("allNodeDataSources")
  private List<DataSource> allNodeDs;

  @Override
  public DataSource dataSource(Object tenantId) {
    if (tenantId == null) {
      return coordTenant.dataSource();
    } else {
      CitusTenant citusTenant = (CitusTenant) tenantId;
      return citusTenant.dataSource();
    }
  }

  @Override
  public void shutdown(boolean deregisterDriver) {
    System.out.println("Shutdown all DataSources");
    allNodeDs.stream()
            .map(t -> (DataSourcePool) t)
            .forEach(ds -> shutdownDataSource(ds, deregisterDriver));
  }

  private void shutdownDataSource(DataSource dataSource, boolean deregisterDriver) {
    if (dataSource instanceof DataSourcePool) {
      ((DataSourcePool) dataSource).shutdown(deregisterDriver);
    }
  }
}
